package com.sjb.cp.manager.sharding;

import com.baomidou.mybatisplus.extension.toolkit.SqlRunner;
import com.sjb.cp.model.enums.SpaceLevelEnum;
import com.sjb.cp.model.enums.SpaceTypeEnum;
import com.sjb.cp.model.po.Space;
import com.sjb.cp.service.SpaceService;
import lombok.extern.slf4j.Slf4j;
import org.apache.shardingsphere.driver.jdbc.core.connection.ShardingSphereConnection;
import org.apache.shardingsphere.infra.config.rule.RuleConfiguration;
import org.apache.shardingsphere.infra.metadata.database.rule.ShardingSphereRuleMetaData;
import org.apache.shardingsphere.mode.manager.ContextManager;
import org.apache.shardingsphere.sharding.api.config.ShardingRuleConfiguration;
import org.apache.shardingsphere.sharding.api.config.rule.ShardingTableRuleConfiguration;
import org.apache.shardingsphere.sharding.rule.ShardingRule;
import org.apache.shardingsphere.sharding.spi.ShardingAlgorithm;
import org.springframework.stereotype.Component;
import org.springframework.web.bind.annotation.PostMapping;

import javax.annotation.PostConstruct;
import javax.annotation.Resource;
import javax.sql.DataSource;
import java.sql.SQLException;
import java.util.Collections;
import java.util.List;
import java.util.Optional;
import java.util.Set;
import java.util.stream.Collectors;

/**
 * @Author:sjb
 * @CreateTime:2025-02-19
 * @Description: 动态分表管理器
 * @Version：1.0
 */
@Slf4j
@Component
public class DynamicShardingManager {


    private static final String LOGIC_TABLE_NAME = "picture";

    // 配置文件中的数据库名称
    private static final String DATABASE_NAME = "logic_db";

    @Resource
    private DataSource dataSource;

    @Resource
    private SpaceService spaceService;


    @PostConstruct
    public void init() {
        log.info("初始化表配置....");
        // todo
    }

    /**
     * 获取所有动态表名 包括初始表和分表
     *
     * @return
     */
    private Set<String> fetchAllTableNames() {
        Set<Long> spaceIds = spaceService.lambdaQuery()
                .eq(Space::getSpaceType, SpaceTypeEnum.TEAM.getValue())
                .list().stream().map(Space::getId).collect(Collectors.toSet());
        Set<String> tableNames = spaceIds.stream()
                .map(spaceId -> LOGIC_TABLE_NAME + "_" + spaceId).collect(Collectors.toSet());
        tableNames.add(LOGIC_TABLE_NAME);
        return tableNames;
    }

    /**
     * 更新 ShardingSphere 的 actual-data-nodes 动态表名配置
     */
    private void updateShardingTableNodes() {
        Set<String> taleNames = fetchAllTableNames();
        String newActualDataNodes = taleNames.stream().map(tableName -> "sjb_picture." + tableName)
                .collect(Collectors.joining(","));
        log.info("动态分表actual-data-nodes 配置: {}", newActualDataNodes);

        ContextManager contextManager = getContextManager();
        ShardingSphereRuleMetaData ruleMetaData = contextManager.getMetaDataContexts()
                .getMetaData()
                .getDatabases()
                .get(DATABASE_NAME)
                .getRuleMetaData();

        Optional<ShardingRule> singleRule = ruleMetaData.findSingleRule(ShardingRule.class);
        if (singleRule.isPresent()) {
            ShardingRuleConfiguration ruleConfiguration = (ShardingRuleConfiguration) singleRule.get().getConfiguration();
            List<ShardingTableRuleConfiguration> updateRules = ruleConfiguration.getTables().stream().map(
                    table -> {
                        if (table.getLogicTable().equals(LOGIC_TABLE_NAME)) {
                            ShardingTableRuleConfiguration newTableRuleConfig = new ShardingTableRuleConfiguration(LOGIC_TABLE_NAME, newActualDataNodes);
                            newTableRuleConfig.setDatabaseShardingStrategy(table.getDatabaseShardingStrategy());
                            newTableRuleConfig.setTableShardingStrategy(table.getTableShardingStrategy());
                            newTableRuleConfig.setKeyGenerateStrategy(table.getKeyGenerateStrategy());
                            newTableRuleConfig.setAuditStrategy(table.getAuditStrategy());
                            return newTableRuleConfig;
                        }
                        return table;
                    }).collect(Collectors.toList());
            ruleConfiguration.setTables(updateRules);
            contextManager.alterRuleConfiguration(DATABASE_NAME, Collections.singleton(ruleConfiguration));
            contextManager.reloadDatabase(DATABASE_NAME);
            log.info("动态分表规则更新成功！");
        }else{
            log.error("未找到 ShardingSphere 的分片规则配置，动态分表更新失败。");
        }

    }


    /**
     * 获取动态分表的contextManager
     *
     * @return
     */
    private ContextManager getContextManager() {
        try (ShardingSphereConnection connection = dataSource.getConnection().unwrap(ShardingSphereConnection.class)) {
            return connection.getContextManager();
        } catch (SQLException exception) {
            throw new RuntimeException("获取 ShardingSphere ContextManager 失败", exception);
        }
    }


    /**
     * 创建动态分表
     * @param space
     */
    public void createSpacePictureTable(Space space){
        if (space.getSpaceType().equals(SpaceTypeEnum.TEAM.getValue())
                && space.getSpaceLevel().equals(SpaceLevelEnum.FLAGSHIP.getValue())){
            Long spaceId = space.getId();
            String tableName = "picture_" + spaceId;

            String createSql = "create table " + tableName + "LIKE picture";

            try {
                SqlRunner.db().update(createSql);
                // 更新分表
                updateShardingTableNodes();
            } catch (Exception e) {
                log.error("创建图片空间分表失败，空间 id = {}", space.getId());
            }
        }
    }

}
